from Model_Simulation import * 
import numpy as np 
from tensorforce.agents import Agent
from utils_agent import * 
from Parameters_1D import *
totalDepth, axialNodes, totalNodes = spatialVariables_1D_act() 

LAI_factors = np.loadtxt(".../relevant_data/LAI_factors_reduced.txt")

#Load the agent
agent_name = 'Scheduler_agent -4-4'
agent_loaded = Agent.load("./Trained_Agent/",filename=agent_name)
internals = agent_loaded.initial_internals()

#Define some simulation metrics and relevant arrays
episode_len = 14 
stateArray = np.zeros((episode_len+1, axialNodes+4))
stateArray_rz = np.zeros(episode_len+1)
actionsArray = np.zeros(episode_len)

#Define the bounds of the zone
lower_zone = 0.176
upper_zone = 0.280

def compute_rootZone_moisture(current_state):
        if current_state[-2] == 0.50:
            rootZone_vol_moist = 0.10*((1/6)*(current_state[10] + current_state[11] + current_state[12] + current_state[13] + current_state[14] + current_state[15])) +\
                         0.20*((1/6)*(current_state[15] + current_state[16] + current_state[17] + current_state[18] + current_state[19] + current_state[20])) +\
                         0.30*((1/6)*(current_state[20] + current_state[21] + current_state[22] + current_state[23] + current_state[24] + current_state[25])) + \
                         0.40*((1/6)*(current_state[25] + current_state[26] + current_state[27] + current_state[28] + current_state[29] + current_state[30]))
        else:
            rootZone_vol_moist = 0.10*((1/6)*(current_state[0] + current_state[1] + current_state[2] + current_state[3] + current_state[4] + current_state[5]))  + \
                         0.20*((1/6)*(current_state[5] + current_state[6] + current_state[7] + current_state[8] + current_state[9] + current_state[10])) + \
                         0.30*((1/11)*(current_state[10] + current_state[11] + current_state[12] + current_state[13] + current_state[14] + current_state[15] + current_state[16] + current_state[17] + current_state[18] + current_state[19] + current_state[20])) + \
                         0.40*((1/11)*(current_state[20] + current_state[21] + current_state[22] + current_state[23] + current_state[24] + current_state[25] + current_state[26] + current_state[27] + current_state[28] + current_state[29] + current_state[30]))

        return rootZone_vol_moist

#Set the initial state
current_rootDepth = 1.00
rooting_depths = np.zeros(episode_len+1)
rooting_depths[0:15] = 0.50
rooting_depths[15:] = 1.00

lai_factors = LAI_factors[0:episode_len+1]

state_init = np.zeros(axialNodes+4)
state_init[0:axialNodes] = 0.22*np.ones(axialNodes)
state_init[axialNodes]   =  np.random.uniform(1.04, 9.0)
state_init[axialNodes+1] = np.random.uniform(0.20, 1.25)
state_init[axialNodes+2] = rooting_depths[0]
state_init[axialNodes+3] = LAI_factors[0]
stateArray_rz[0] = compute_rootZone_moisture(state_init)
stateArray[0,:] = state_init

for i in range(episode_len):
    #Generate some noise and add same to the states and weather data
    # noise_states = 0*np.random.uniform(0.01, 0.05, size=(axialNodes)) # Using a uniform distribution
    # noise_weather = np.zeros(2)
    # noise_all = np.concatenate((noise_states,noise_weather))
    # states = stateArray[i,:] + noise_all
    states = stateArray[i,:] 

    for j in range(axialNodes):
        if states[j] < 0.078:
            states[j] = 0.078
            pass
        if states[j] > 0.43:
            states[j] = 0.43
            pass
    actions, internals = agent_loaded.act(states=states, internals=internals, independent=True)
    actionsArray[i] = -1*actions["irrigationRate"]*actions["irrigationDecision"]
    current_state=stateArray[i,:]
    next_state = SimulateModel(current_state[0:axialNodes], actions["irrigationRate"]*actions["irrigationDecision"],
                               current_state[axialNodes], current_state[axialNodes+1], current_state[-2], current_state[-1])
    
    next
    
    #next_state[-1] = current_rootDepth
    next_state[-2] = rooting_depths[i+1]
    next_state[-1] = lai_factors[i+1]
    stateArray_rz[i+1] = compute_rootZone_moisture(next_state)
    stateArray[i+1,:] = next_state
    pass 

#Generate and save a plot
filename = 'Trained_Agent_Res_3'
GenerateTrajectories(episode_len, actionsArray,stateArray_rz,lower_zone, upper_zone,filename)

# #Generate a plot of the reward trajectory
# filename_reward = 'Reward_Trajectory_5_mz1'
# reward_trajectory = np.loadtxt("reward_trajectory_5.txt")
# PlotRewardTrajectory(rewardTrajectory=reward_trajectory, filename=filename_reward)
